package branch_and_price;

import gurobi.GRB;
import gurobi.GRBEnv;
import gurobi.GRBException;
import gurobi.GRBLinExpr;
import gurobi.GRBModel;
import gurobi.GRBVar;

import java.util.Arrays;

import java.util.ArrayList;

public class Solver {

    GRBEnv env;
    GRBModel rmp;

    private int m;
    private int n;

    private ArrayList<ArrayList<double[]>> assignments;
    private ArrayList<int[]> fix_zero = new ArrayList<int[]>();
    private ArrayList<int[]> fix_one = new ArrayList<int[]>();
    private int k;

    // 当前模型的解及对偶
    private double obj = -Double.MAX_VALUE;
    private double[] variables;
    private double[] dual;

    // 目前为止的最优解（整数解）
    private double optimal = -Double.MAX_VALUE;
    private double[] optVars;

    Solver(ArrayList<ArrayList<double[]>> assignments) {
        this.m = assignments.size();
        this.n = assignments.get(0).get(0).length;
        // 统计方案总数
        this.k = 0;
        for (int i = 0; i < m; i++) {
            this.k += assignments.get(i).size();
        }
        this.assignments = assignments;
        try {
            this.env = new GRBEnv(true);
            env.start();
            this.rmp = new GRBModel(env);
            this.rmp.set("OutputFlag", "0");
        } catch (GRBException e) {
            System.out.println("Error code: " + e.getErrorCode() + ". " + e.getMessage());
        }
    }

    // 根据assignments生成RMP
    private void constructRMP() {
        try {
            rmp = new GRBModel(env);
            rmp.set("OutputFlag", "0");

            int numOfVars = 0;
            for (int i = 0; i < m; i++) {
                numOfVars += assignments.get(i).size();
            }
            GRBVar[] y = rmp.addVars(numOfVars, GRB.CONTINUOUS);
            GRBVar z = rmp.addVar(0, 1, 0, GRB.CONTINUOUS, "dummy");

            double[] rmp_c = new double[numOfVars];
            int index = 0;
            for (int i = 0; i < m; i++) {
                for (int kk = 0; kk < assignments.get(i).size(); kk++) {
                    double sum = 0;
                    for (int j = 0; j < n; j++) {
                        sum += Dataset.profit[i][j] * assignments.get(i).get(kk)[j];
                    }
                    rmp_c[index++] = sum;
                }
            }
            GRBLinExpr expr = new GRBLinExpr();
            expr.addTerms(rmp_c, y);
            expr.addTerm(-1, z);
            rmp.setObjective(expr, GRB.MAXIMIZE);

            for (int j = 0; j < n; j++) {
                index = 0;
                expr = new GRBLinExpr();
                double[] rmp_constr = new double[numOfVars];
                for (int i = 0; i < m; i++) {
                    for (int kk = 0; kk < assignments.get(i).size(); kk++) {
                        rmp_constr[index++] = assignments.get(i).get(kk)[j];
                    }
                }
                expr.addTerms(rmp_constr, y);
                expr.addTerm(1, z);
                rmp.addConstr(expr, GRB.EQUAL, 1, "job " + j);
            }

            index = 0;
            for (int i = 0; i < m; i++) {
                expr = new GRBLinExpr();
                int temp = assignments.get(i).size();
                double[] rmp_constr = new double[temp];
                for (int j = 0; j < temp; j++) {
                    rmp_constr[j] = 1;
                }
                expr.addTerms(rmp_constr, Arrays.copyOfRange(y, index, index + temp));
                expr.addTerm(1, z);
                rmp.addConstr(expr, GRB.LESS_EQUAL, 1, "agent " + i);
                index += temp;
            }

            // 如果有需要fix的变量，则加入对应的约束
            for (int[] v : fix_zero) {
                int sum = 0;
                for (int i = 0; i < m; i++) {
                    if (i != v[0]) {
                        // 计数，因为变量y的下标一维化了，只靠i和j无法直接定位y
                        sum += assignments.get(i).size();
                    } else {
                        for (int k = 0; k < assignments.get(i).size(); k++) {
                            if (assignments.get(i).get(k)[v[1]] != 0) {
                                expr = new GRBLinExpr();
                                expr.addTerm(1, y[sum + k]);
                                rmp.addConstr(expr, GRB.EQUAL, 0, "fix 0");
                            }
                        }
                        break;
                    }
                }
            }
            for (int[] v : fix_one) {
                int sum = 0;
                for (int i = 0; i < m; i++) {
                    if (i != v[0]) {
                        // 计数，因为变量y的下标一维化了，只靠i和j无法直接定位y
                        sum += assignments.get(i).size();
                    } else {
                        for (int k = 0; k < assignments.get(i).size(); k++) {
                            if (assignments.get(i).get(k)[v[1]] != 1) {
                                expr = new GRBLinExpr();
                                expr.addTerm(1, y[sum + k]);
                                rmp.addConstr(expr, GRB.EQUAL, 0, "fix 1");
                            }
                        }
                        break;
                    }
                }
            }

            // rmp.write("1.lp");

        } catch (GRBException e) {
            System.out.println("Error code: " + e.getErrorCode() + ". " + e.getMessage());
        }

        // System.out.println("当前方案为：");
        // for (int i = 0; i < m; i++) {
        // System.out.println("机器" + i);
        // for (double[] v : assignments.get(i)) {
        // System.out.println(Arrays.toString(v));
        // }
        // }
    }

    // 构造pricing problem并求解，返回是否可以改进RMP
    private boolean solvePricing() {
        int agent = -1;
        double maximum = 0;
        double[] solution = new double[n];
        // 求解i个pricing problem
        for (int i = 0; i < m; i++) {
            try {

                GRBModel pricing = new GRBModel(env);
                pricing.set("OutputFlag", "0");

                GRBVar[] x = pricing.addVars(n, GRB.BINARY);

                GRBLinExpr expr = new GRBLinExpr();
                double[] pricing_c = Utils.p_i(i);
                for (int j = 0; j < n; j++) {
                    pricing_c[j] -= dual[j];
                }
                expr.addTerms(pricing_c, x);
                pricing.setObjective(expr, GRB.MAXIMIZE);

                expr = new GRBLinExpr();
                expr.addTerms(Utils.w_i(i), x);
                pricing.addConstr(expr, GRB.LESS_EQUAL, Utils.c(i), "c0");

                // 加入fix变量对应的约束
                for (int[] v : fix_zero) {
                    if (v[0] == i) {
                        expr = new GRBLinExpr();
                        expr.addTerm(1, x[v[1]]);
                        pricing.addConstr(expr, GRB.EQUAL, 0, "fix 0");
                    }
                }
                for (int[] v : fix_one) {
                    if (v[0] == i) {
                        expr = new GRBLinExpr();
                        expr.addTerm(1, x[v[1]]);
                        pricing.addConstr(expr, GRB.EQUAL, 1, "fix 1");
                    }
                }

                pricing.optimize();

                if (maximum < pricing.get(GRB.DoubleAttr.ObjVal) - dual[n + i]) {
                    agent = i;
                    maximum = pricing.get(GRB.DoubleAttr.ObjVal) - dual[n + i];
                    for (int j = 0; j < n; j++) {
                        solution[j] = pricing.getVars()[j].get(GRB.DoubleAttr.X);
                    }
                }

            } catch (GRBException e) {
                System.out.println("Error code: " + e.getErrorCode() + ". " + e.getMessage());
            }
        }
        // 更新对应方案
        if (maximum > 0) {
            assignments.get(agent).add(solution);
            k++;
            return true;
        } else {
            return false;
        }
    }

    // 更新当前解
    private void updateCurrentSolution() {
        try {
            obj = rmp.get(GRB.DoubleAttr.ObjVal);
            variables = new double[rmp.getVars().length - 1];
            for (int i = 0; i < rmp.getVars().length - 1; i++) {
                variables[i] = rmp.getVars()[i].get(GRB.DoubleAttr.X);
            }
        } catch (GRBException e) {
            System.out.println("Error code: " + e.getErrorCode() + ". " + e.getMessage());
        }
    }

    // 更新最优解
    private void updateOptimal() {
        optimal = obj;
        optVars = Arrays.copyOf(variables, variables.length);
    }

    // 判断变量是否都为整数
    private boolean intVariables(GRBVar[] vars) {
        for (GRBVar i : vars) {
            try {
                double v = i.get(GRB.DoubleAttr.X);
                if (v != (int) v) {
                    return false;
                }
            } catch (GRBException e) {
                System.out.println("Error code: " + e.getErrorCode() + ". " + e.getMessage());
            }
        }
        return true;
    }

    // 获取需要分支的变量，返回i和j
    private int[] getBranchingVariable() {
        int index = 0;
        double v = 0.5;
        int bi = -1;
        int bj = -1;
        for (int i = 0; i < m; i++) {
            for (int j = 0; j < n; j++) {
                double sum = 0;
                for (int k = 0; k < assignments.get(i).size(); k++) {
                    sum += assignments.get(i).get(k)[j] * variables[index + k];
                }
                if (Utils.abs(sum - 0.5) < v) {
                    v = Utils.abs(sum - 0.5);
                    bi = i;
                    bj = j;
                }
            }
            index += assignments.get(i).size();
        }
        return new int[] { bi, bj };
    }

    // 复制assignments
    private ArrayList<ArrayList<double[]>> copyAssignments(ArrayList<ArrayList<double[]>> a) {
        ArrayList<ArrayList<double[]>> b = new ArrayList<ArrayList<double[]>>(m);
        for (int i = 0; i < m; i++) {
            ArrayList<double[]> temp = new ArrayList<double[]>();
            for (int j = 0; j < a.get(i).size(); j++) {
                temp.add(Arrays.copyOf(a.get(i).get(j), n));
            }
            b.add(temp);
        }
        return b;
    }

    // 复制fix
    private ArrayList<int[]> copyFix(ArrayList<int[]> f) {
        ArrayList<int[]> b = new ArrayList<int[]>(f.size());
        for (int i = 0; i < f.size(); i++) {
            b.add(Arrays.copyOf(f.get(i), 2));
        }
        return b;
    }

    public void solve() {

        // 构造初始RMP
        constructRMP();

        // 开始求解
        solveProblem();

        // 输出最优解
        System.out.println(optimal);
        System.out.println(Arrays.toString(optVars));

    }

    public void solveProblem() {
        while (true) {
            try {
                rmp.optimize();
                dual = rmp.get(GRB.DoubleAttr.Pi, rmp.getConstrs());
                boolean improved = solvePricing();
                if (!improved) {
                    if (optimal < rmp.get(GRB.DoubleAttr.ObjVal)) {
                        updateCurrentSolution();
                        if (!intVariables(rmp.getVars())) {
                            System.out.println("开始分支");
                            branching();
                        } else {
                            updateOptimal();
                            System.out.println("找到一组整数解");
                        }
                        break;
                    } else {
                        System.out.println("剪枝");
                        break;
                    }
                } else {
                    constructRMP();
                }
            } catch (GRBException e) {
                System.out.println("Error code: " + e.getErrorCode() + ". " + e.getMessage());
            }
        }
    }

    public void branching() {

        // 获取需要分支的变量
        int[] v = getBranchingVariable();

        // 生成x=0的分支
        ArrayList<ArrayList<double[]>> assignments_0 = copyAssignments(assignments);
        ArrayList<int[]> fix_zero_0 = copyFix(fix_zero);
        ArrayList<int[]> fix_one_0 = copyFix(fix_one);
        fix_zero_0.add(v);
        int k_0 = k;

        // 生成x=1的分支
        ArrayList<ArrayList<double[]>> assignments_1 = copyAssignments(assignments);
        ArrayList<int[]> fix_zero_1 = copyFix(fix_zero);
        ArrayList<int[]> fix_one_1 = copyFix(fix_one);
        fix_one_1.add(v);
        int k_1 = k;

        // 探索x=0的分支
        assignments = assignments_0;
        fix_zero = fix_zero_0;
        fix_one = fix_one_0;
        k = k_0;
        System.out.println("探索0分支");
        constructRMP();
        solveProblem();

        // 探索x=1的分支
        assignments = assignments_1;
        fix_zero = fix_zero_1;
        fix_one = fix_one_1;
        k = k_1;
        System.out.println("探索1分支");
        constructRMP();
        solveProblem();

    }

}
